#include "unitree_arm_sdk/control/unitreeArm.h"
#include <iostream>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <cstring>
#include <iomanip>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <atomic>
#include <fcntl.h>
#include <sys/select.h>
#include <fstream>
#include <chrono>
#include <ctime>
#include <sstream>
#include <string>
#include <cerrno>
#include <termios.h>
#include "unitree_arm_sdk/thirdparty/robotics.h"    // for rpyxyz -> HomoMat
#include "unitree_arm_sdk/model/ArmModel.h"         // includes ArmModel, Z1Model
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <termios.h>
#include <fcntl.h>
#include <unistd.h>
#include <iomanip>
#include <cstring>

// Choose one method. Here we do position control via MoveJ:
#define USE_POSITION_CONTROL
// #define USE_SPEED_CONTROL

// threshold for micro adjust, if you choose to keep that logic
#define thresh 0.2

// Uncomment if you want to restore predictive logic (not used in this example).
// #define USE_PREDICTIVE_CONTROL

using namespace UNITREE_ARM;

/* ====================== Global variables / Functions for Force Reading ====================== */

// We store the "latest packet" from TCP in a single global variable:
static std::mutex g_latestMutex;
static std::vector<unsigned char> g_latestPacket;
static std::atomic<bool> g_hasNewPacket(false);

// Global force data
float force[6] = {0.0f}; // store the last 6 channels
std::mutex force_mutex;
std::condition_variable force_cv;

// Global exit flag
std::atomic<bool> exit_flag(false);
std::mutex arm_mutex;

// CRC helper
uint16_t crc16_modbus(const uint8_t *data, size_t len) {
    uint16_t crc = 0xFFFF;
    for (size_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i];
        for (int j = 0; j < 8; j++) {
            if (crc & 0x0001) {
                crc >>= 1;
                crc ^= 0xA001;
            } else {
                crc >>= 1;
            }
        }
    }
    return crc;
}

// Debug printing
static void print_header_check(const unsigned char* response, size_t len) {
    if (len >= 2) {
        std::cout << "First 2 bytes: ";
        for (int i = 0; i < 2; i++) {
            std::cout << std::hex << std::uppercase 
                      << std::setw(2) << std::setfill('0') << (int)response[i] << " ";
        }
        std::cout << "(Hex)" << std::dec << std::endl;
        std::cout << "Response Length: " << len << std::endl;
    }
}

// Parse response from the sensor
void parse_response(const unsigned char* response, size_t len) {
    if (len < 27) {
        std::cout << "Invalid response length." << std::endl;
        return;
    }
    print_header_check(response, len);  // Print first two bytes for debugging

    // Check header
    if (response[0] != 0xAA || response[1] != 0x55) {
        std::cout << "Invalid header." << std::endl;
        return;
    }
    // Data portion
    float values[6];
    memcpy(values, response + 6, 6 * sizeof(float));
    for (int i = 0; i < 6; i++) {
        force[i] = values[i];
    }
}

/* ====================== Serial Port Setup for Force Sensor ====================== */

int initSerialPort(const std::string& portname) {
    int fd = open(portname.c_str(), O_RDWR | O_NOCTTY | O_SYNC);
    if (fd < 0) {
        std::cerr << "[initSerialPort] Error opening " << portname << ": " 
                  << strerror(errno) << std::endl;
        return -1;
    }

    struct termios tty;
    memset(&tty, 0, sizeof tty);
    if (tcgetattr(fd, &tty) != 0) {
        std::cerr << "[initSerialPort] Error from tcgetattr: " << strerror(errno) << std::endl;
        close(fd);
        return -1;
    }

    cfsetospeed(&tty, B115200);
    cfsetispeed(&tty, B115200);

    tty.c_cflag = (tty.c_cflag & ~CSIZE) | CS8;  
    tty.c_iflag &= ~IGNBRK;        
    tty.c_lflag = 0;               
    tty.c_oflag = 0;               
    tty.c_cc[VMIN]  = 0;           
    tty.c_cc[VTIME] = 5;           

    tty.c_iflag &= ~(IXON | IXOFF | IXANY);
    tty.c_cflag |= (CLOCAL | CREAD);
    tty.c_cflag &= ~(PARENB | PARODD);
    tty.c_cflag &= ~CSTOPB;
    tty.c_cflag &= ~CRTSCTS;

    if (tcsetattr(fd, TCSANOW, &tty) != 0) {
        std::cerr << "[initSerialPort] Error from tcsetattr: " 
                  << strerror(errno) << std::endl;
        close(fd);
        return -1;
    }
    return fd;
}

/* ====================== Custom Arm Class (child of unitreeArm) ====================== */

class Z1ARM : public unitreeArm {
public:
    Z1ARM() : unitreeArm(true), tcp_sockfd(-1) {};
    ~Z1ARM() {
        if (tcp_sockfd != -1) {
            close(tcp_sockfd);
            std::cout << "[~Z1ARM] TCP socket closed." << std::endl;
        }
    }

    // We keep a few example methods from your original
    void armCtrlByFSM();    
    void printState();

    // Additional fields
    double gripper_pos = 0.0;
    double joint_speed = 2.0;
    double cartesian_speed = 3.0;

    // For demonstration, we store one "forward_position"
    void setForwardPosition(const Vec6& pos) { forward_position = pos; }
    Vec6 getForwardPosition() const { return forward_position; }

    int tcp_sockfd;
private:
    Vec6 forward_position; 
};

// Simple example using finite state machine
void Z1ARM::armCtrlByFSM() {
    std::cout << "[armCtrlByFSM] Before labelRun(\"forward\")" << std::endl;
    try {
        labelRun("forward");
        std::cout << "[armCtrlByFSM] labelRun(\"forward\") OK." << std::endl;
    }
    catch (...) {
        std::cerr << "[armCtrlByFSM] Exception in labelRun." << std::endl;
        return;
    }
    printState();
    {
        std::lock_guard<std::mutex> lock(arm_mutex);
        forward_position = lowstate->endPosture;
        std::cout << "[armCtrlByFSM] Save forward_position: " << forward_position.transpose() << std::endl;
    }
}

void Z1ARM::printState() {
    std::cout << "------ Joint State ------" << std::endl;
    std::cout << "qState: " << lowstate->getQ().transpose() << std::endl;
    std::cout << "qdState: " << lowstate->getQd().transpose() << std::endl;
    std::cout << "tauState: " << lowstate->getTau().transpose() << std::endl;
    std::cout << "------ Endeffector Cartesian Posture ------" << std::endl;
    std::cout << "(roll pitch yaw x y z) = " << lowstate->endPosture.transpose() << std::endl;
}

/* ====================== Keyboard Monitor ====================== */
void monitorKeyboard(std::atomic<bool>& exit_flag) {
    std::cout << "[monitorKeyboard] Press 'q' followed by Enter to exit." << std::endl;
    char c;
    while (std::cin >> c) {
        if (c == 'q' || c == 'Q') {
            exit_flag = true;
            std::cout << "[monitorKeyboard] 'q' pressed. Exiting..." << std::endl;
            force_cv.notify_all();
            break;
        }
    }
}

/* ====================== TCP Client: Connect to server ====================== */
void tcpClient(Z1ARM& arm) {
    arm.tcp_sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (arm.tcp_sockfd < 0) {
        std::cerr << "[tcpClient] Failed to create socket: " << strerror(errno) << std::endl;
        return;
    }

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(14532);
    server_addr.sin_addr.s_addr = inet_addr("115.236.153.174");

    std::cout << "[tcpClient] Attempt to connect..." << std::endl;

    // Non-blocking connect with select():
    int flags = fcntl(arm.tcp_sockfd, F_GETFL, 0);
    fcntl(arm.tcp_sockfd, F_SETFL, flags | O_NONBLOCK);

    connect(arm.tcp_sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr));

    fd_set set;
    FD_ZERO(&set);
    FD_SET(arm.tcp_sockfd, &set);

    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;

    int ret = select(arm.tcp_sockfd + 1, NULL, &set, NULL, &timeout);
    if (ret <= 0) {
        std::cerr << "[tcpClient] Connect timed out or error: " 
                  << strerror(errno) << std::endl;
        close(arm.tcp_sockfd);
        arm.tcp_sockfd = -1;
        return;
    } else {
        int so_error = 0;
        socklen_t len = sizeof(so_error);
        getsockopt(arm.tcp_sockfd, SOL_SOCKET, SO_ERROR, &so_error, &len);
        if (so_error != 0) {
            std::cerr << "[tcpClient] Connect failed. so_error="<< so_error 
                      << ", " << strerror(so_error) << std::endl;
            close(arm.tcp_sockfd);
            arm.tcp_sockfd = -1;
            return;
        }
    }

    // If we get here, we connected
    std::cout << "[tcpClient] Connected OK." << std::endl;

    // Return the socket in arm.tcp_sockfd
}

/* ====================== Send Force Data to Server ====================== */
void sendForceData(Z1ARM& arm, float* force) {
    if (arm.tcp_sockfd == -1) {
        std::cerr << "[sendForceData] TCP not connected." << std::endl;
        return;
    }

    // Packet layout:  4 header bytes + (4 bytes * 6 floats) + 2 CRC = 4 + 24 + 2 = 30 bytes
    uint8_t packet[4 + 4 * 6 + 2];
    packet[0] = 0xAA;
    packet[1] = 0xBB;
    packet[2] = 0xCC;
    packet[3] = 0xDD;

    // Copy 6 floats of force
    std::memcpy(&packet[4], force, 6 * sizeof(float));

    // Compute CRC on the first (4+24)=28 bytes
    uint16_t crc = crc16_modbus(packet, 4 + (6 * 4));
    packet[28] = (uint8_t)(crc & 0xFF);
    packet[29] = (uint8_t)((crc >> 8) & 0xFF);

    // Send
    ssize_t sent = send(arm.tcp_sockfd, (char*)packet, sizeof(packet), 0);
    if (sent < 0) {
        std::cerr << "[sendForceData] Failed: " << strerror(errno) << std::endl;
    }
}

/* ====================== Logging Arm Position ====================== */
void logArmPosition(Z1ARM& arm, std::atomic<bool>& exit_flag, const std::string& filename) {
    std::ofstream outfile(filename, std::ios::out);
    if (!outfile.is_open()) {
        std::cerr << "[logArmPosition] Cannot open file: " << filename << std::endl;
        return;
    }

    std::cout << "[logArmPosition] Logging posture to " << filename << std::endl;
    outfile << "timestamp_ms roll pitch yaw x y z Fx Fy Fz Tx Ty Tz "
            << "q0 q1 q2 q3 q4 q5 qd0 qd1 qd2 qd3 qd4 qd5 tau0 tau1 tau2 tau3 tau4 tau5\n";
    outfile.flush();

    std::chrono::milliseconds interval(10);
    while (!exit_flag) {
        auto now = std::chrono::system_clock::now();
        auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                          now.time_since_epoch()) 
                      % 1000;
        auto now_time_t = std::chrono::system_clock::to_time_t(now);
        std::tm tm_time;
        localtime_r(&now_time_t, &tm_time);

        std::ostringstream timestamp_stream;
        timestamp_stream << std::put_time(&tm_time, "%Y-%m-%d_%H-%M-%S") 
                         << '.' << std::setfill('0') << std::setw(3) 
                         << now_ms.count();

        Vec6 endPosture;
        {
            std::lock_guard<std::mutex> lk(arm_mutex);
            if (arm.lowstate == nullptr) {
                endPosture = Vec6::Zero();
            } else {
                endPosture = arm.lowstate->endPosture;
            }
        }

        Vec6 Q  = arm.lowstate->getQ();
        Vec6 Qd = arm.lowstate->getQd();
        Vec6 Qt = arm.lowstate->getTau();

        float force_vals[6];
        {
            std::lock_guard<std::mutex> lk(force_mutex);
            std::copy(std::begin(force), std::end(force), std::begin(force_vals));
        }

        // Write a line of data
        outfile << timestamp_stream.str() << " "
                << endPosture[0] << " " << endPosture[1] << " " << endPosture[2] << " "
                << endPosture[3] << " " << endPosture[4] << " " << endPosture[5] << " "
                << force_vals[0] << " " << force_vals[1] << " " << force_vals[2] << " "
                << force_vals[3] << " " << force_vals[4] << " " << force_vals[5] << " "
                << Q[0] << " " << Q[1] << " " << Q[2] << " " << Q[3] << " " << Q[4] << " " << Q[5] << " "
                << Qd[0] << " " << Qd[1] << " " << Qd[2] << " " << Qd[3] << " " << Qd[4] << " " << Qd[5] << " "
                << Qt[0] << " " << Qt[1] << " " << Qt[2] << " " << Qt[3] << " " << Qt[4] << " " << Qt[5]
                << "\n";
        outfile.flush();

        std::this_thread::sleep_for(interval);
    }

    outfile.close();
    std::cout << "[logArmPosition] Logging stopped.\n";
}

/* ====================== TCP Receive: Store only the latest packet ====================== */
void receiveDataFromServer(Z1ARM& arm) {
    // We'll parse 40 bytes or 44, etc., depending on your packet structure
    // In your code, the final version used 40 or 44. Let’s keep 40 for the minimal example.

    const int PACKET_SIZE = 40;
    std::vector<unsigned char> buffer;
    // Make socket non-blocking
    int flags = fcntl(arm.tcp_sockfd, F_GETFL, 0);
    if (flags == -1) {
        std::cerr << "[receiveDataFromServer] fcntl get err.\n";
        exit_flag = true;
        return;
    }
    if (fcntl(arm.tcp_sockfd, F_SETFL, flags | O_NONBLOCK) == -1) {
        std::cerr << "[receiveDataFromServer] fcntl set O_NONBLOCK err.\n";
        exit_flag = true;
        return;
    }

    while (!exit_flag) {
        unsigned char tmp[512];
        ssize_t bytes = recv(arm.tcp_sockfd, (char*)tmp, sizeof(tmp), 0);
        if (bytes < 0) {
            if (errno == EWOULDBLOCK || errno == EAGAIN) {
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            } else {
                std::cerr << "[receiveDataFromServer] recv err: " 
                          << strerror(errno) << "\n";
                exit_flag = true;
                break;
            }
        } else if (bytes == 0) {
            std::cerr << "[receiveDataFromServer] remote closed.\n";
            exit_flag = true;
            break;
        }

        // Accumulate data in buffer
        buffer.insert(buffer.end(), tmp, tmp + bytes);

        // Attempt to extract complete frames
        while (buffer.size() >= PACKET_SIZE && !exit_flag) {
            size_t pos = buffer.size();
            bool found_header = false;
            // Look for 4-byte header AA BB CC DD
            for (size_t i = 0; i + PACKET_SIZE <= buffer.size(); i++) {
                if (buffer[i]   == 0xAA && buffer[i+1] == 0xBB &&
                    buffer[i+2] == 0xCC && buffer[i+3] == 0xDD) {
                    pos = i;
                    found_header = true;
                    break;
                }
            }
            if (!found_header) {
                // no valid header, discard all
                buffer.clear();
                break;
            }
            if (buffer.size() - pos >= PACKET_SIZE) {
                std::vector<unsigned char> packet(
                    buffer.begin() + pos,
                    buffer.begin() + pos + PACKET_SIZE);
                // For illustration, let's just check last 2 bytes are EE FF:
                if (packet[38] == 0xEE && packet[39] == 0xFF) {
                    // We treat this as a valid packet
                    {
                        std::lock_guard<std::mutex> lk(g_latestMutex);
                        g_latestPacket = packet;
                        g_hasNewPacket.store(true);
                    }
                    buffer.erase(buffer.begin(), buffer.begin() + pos + PACKET_SIZE);
                } else {
                    // Not a valid tail, remove header
                    buffer.erase(buffer.begin(), buffer.begin() + pos + 4);
                }
            } else {
                // incomplete
                break;
            }
        }
    }
}

/* ====================== Actual Arm Control Thread Using MoveJ ====================== */
void armControlThreadFunc(Z1ARM& arm) 
{
    // We’ll do a straightforward approach: whenever a new packet arrives,
    // parse the posture, then call MoveJ() to reach that posture in joint space.

    
    // Optionally make sure we are not “PASSIVE”:
    arm.setWait(false);
    arm.setFsm(ArmFSMState::JOINTCTRL);

    // We can set waiting or not:
    // By default, the Unitree MoveJ() blocks until the motion finishes.
    // If you prefer asynchronous, call: arm.setWait(false);

    while (!exit_flag) {
        if (!g_hasNewPacket.load()) {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            continue;
        }

        // 1) Extract the latest packet
        std::vector<unsigned char> raw_packet;
        {
            std::lock_guard<std::mutex> lk(g_latestMutex);
            raw_packet = g_latestPacket;
            g_hasNewPacket.store(false);
        }
        if (raw_packet.size() < 40) {
            continue; // skip
        }

        // 2) Suppose you store float data at [4..(some offset)]
        float pos_array[7];
        memcpy(pos_array, &raw_packet[4], 7 * sizeof(float));

        // For example, interpret them as x, y, z, pitch, yaw, roll, etc. 
        float x    = pos_array[0];
        float y    = pos_array[1];
        float z    = pos_array[2];
        float pitch= pos_array[3];
        float yaw  = pos_array[4];
        float roll = pos_array[5];
        float micro= pos_array[6]; // e.g. threshold check

        // 3) Construct a posture in the format (roll, pitch, yaw, x, y, z)
        Vec6 new_position;
        new_position[0] = roll;
        new_position[1] = pitch;
        new_position[2] = yaw;
        new_position[3] = 4*z / 1000.0f;  // scale example
        new_position[4] = 4*x / 1000.0f; 
        new_position[5] = 0.31f + 2*y / 1000.0f; 

        std::cout << "[armControlThreadFunc] MoveJ target: " 
                  << new_position.transpose() << std::endl;

        // 4) Call MoveJ
        double maxJointSpeed = 3;  // rad/s
        bool ok = arm.MoveJ(new_position, maxJointSpeed);
        if (!ok) {
            std::cerr << "[armControlThreadFunc] MoveJ() failed (IK or other)." << std::endl;
        } else {
            std::cout << "[armControlThreadFunc] MoveJ done." << std::endl;
            arm.printState();
        }
    }
}

/* ====================== Reading Force from Serial ====================== */
unsigned char cmd[] = "AT+GOD\r\n";
unsigned char response[100]; // Buffer for the force sensor's response

void readSerialData(int fd, std::atomic<bool>& exit_flag) {
    while (!exit_flag) {
        // Typically, flush and send a command to the sensor
        tcflush(fd, TCIFLUSH);
        write(fd, cmd, sizeof(cmd) - 1);
        usleep(1000000); // wait 1 second before reading

        int bytes_read = read(fd, response, sizeof(response));
        if (bytes_read > 0) {
            print_header_check(response, bytes_read);
            if (response[0] == 0xAA && response[1] == 0x55) {
                parse_response(response, bytes_read);
            } else {
                std::cout << "Invalid header from force sensor." << std::endl;
            }
        } else {
            std::cout << "No data received from force sensor." << std::endl;
        }
        usleep(1000000); // another 1 second
    }
    close(fd);
}

/* ====================== Send Force Thread ====================== */
void sendForceThreadFunc(Z1ARM& arm) {
    while (!exit_flag) {
        float force_vals[6];
        {
            std::unique_lock<std::mutex> lk(force_mutex);
            std::copy(std::begin(force), std::end(force), std::begin(force_vals));
        }
        // Debug print if you like:
        // std::cout << "[sendForceThreadFunc] Force data: ";
        // for (int i = 0; i < 6; i++) std::cout << force_vals[i] << " ";
        // std::cout << std::endl;

        sendForceData(arm, force_vals);
        usleep(100000); // 100 ms
    }
}

/* ====================== MAIN ====================== */

int main(){
    std::cout << std::fixed << std::setprecision(3);

    // 1) Open the serial port for the force sensor:
    int serial_fd = initSerialPort("/dev/ttyUSB0");

    // 2) Create arm object and start communication
    Z1ARM arm;
    if(arm.sendRecvThread != nullptr){
        arm.sendRecvThread->start();
        std::cout << "[main] arm comm thread started.\n";
    } else {
        std::cerr << "[main] sendRecvThread=null.\n";
    }

    // 3) Move the arm to start or some known posture
    arm.backToStart();
    arm.labelRun("forward");

    // 4) Start the TCP client connect
    std::thread tcp_thread(tcpClient, std::ref(arm));
    std::cout << "[main] tcp client thread started.\n";

    // 5) Keyboard monitor
    std::thread keyboard_thread(monitorKeyboard, std::ref(exit_flag));
    std::cout << "[main] keyboard monitor thread.\n";

    tcp_thread.join();
    std::cout << "[main] tcp thread joined.\n";

    // 6) Start a logging thread
    auto t    = std::time(nullptr);
    auto tmptr= std::localtime(&t);
    char buf[128];
    if(std::strftime(buf, sizeof(buf), "arm_position_log_%Y%m%d_%H%M%S.txt", tmptr) == 0){
        std::cerr << "[main] strftime fail.\n";
        exit_flag = true;
    }
    std::string filename(buf);
    std::cout << "[main] log file = " << filename << "\n";

    std::thread log_thread;
    if(!exit_flag) {
        log_thread = std::thread(logArmPosition, std::ref(arm), std::ref(exit_flag), filename);
        std::cout << "[main] logging thread.\n";
    }

    // 7) Additional threads
    std::thread receive_thread;
    std::thread control_thread;
    std::thread serial_thread;
    std::thread force_send_thread;

    // If connected successfully:
    if (arm.tcp_sockfd != -1) {
        // optional: small demonstration
        std::thread arm_ctrl_thread(&Z1ARM::armCtrlByFSM, &arm);

        // receiving data from server
        receive_thread = std::thread(receiveDataFromServer, std::ref(arm));

        // reading force sensor from serial
        if (serial_fd >= 0) {
            serial_thread = std::thread(readSerialData, serial_fd, std::ref(exit_flag));
        }

        // sending force data to server
        force_send_thread = std::thread(sendForceThreadFunc, std::ref(arm));

    #ifdef USE_POSITION_CONTROL
        // main control loop (MoveJ in a loop)
        control_thread = std::thread(armControlThreadFunc, std::ref(arm));
        std::cout << "[main] position control thread.\n";
    #elif defined(USE_SPEED_CONTROL)
        // if you wanted to do speed-based control, you would do armSpeedControlThreadFunc
        // control_thread = std::thread(armSpeedControlThreadFunc, std::ref(arm));
    #endif

        // Wait until user hits 'q' or an error
        while (!exit_flag) {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
        std::cout << "[main] exit_flag set, shutting down...\n";

        close(arm.tcp_sockfd);
        arm.tcp_sockfd = -1;
        std::cout << "[main] tcp closed.\n";

        if (receive_thread.joinable()) {
            receive_thread.join();
            std::cout << "[main] recv joined.\n";
        }
        if (control_thread.joinable()) {
            control_thread.join();
            std::cout << "[main] control joined.\n";
        }
        if (arm_ctrl_thread.joinable()) {
            arm_ctrl_thread.join();
            std::cout << "[main] armCtrlByFSM joined.\n";
        }

    } else {
        std::cerr << "[main] cannot connect to server.\n";
        while (!exit_flag) {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }

    force_cv.notify_all();
    if (force_send_thread.joinable()) {
        force_send_thread.join();
        std::cout << "[main] force_send joined.\n";
    }
    if (keyboard_thread.joinable()) {
        keyboard_thread.join();
        std::cout << "[main] kb joined.\n";
    }
    if (log_thread.joinable()) {
        log_thread.join();
        std::cout << "[main] logging joined.\n";
    }
    if (serial_thread.joinable()) {
        serial_thread.join();
        std::cout << "[main] serial joined.\n";
    }

    // 8) Restore arm to a safe posture
    std::cout << "[main] restore arm.\n";
    arm.backToStart();
    arm.setFsm(ArmFSMState::PASSIVE);
    std::cout << "[main] arm restored.\n";

    // 9) Shut down comm thread
    if (arm.sendRecvThread) {
        arm.sendRecvThread->shutdown();
        std::cout << "[main] arm comm thread shut.\n";
    }

    std::cout << "[main] Done.\n";
    return 0;
}
